import asyncio
import msgpack
import time
import traceback
from asyncio.exceptions import TimeoutError
import struct
from lib.log import Log
from lib import capture_error
from bson.objectid import ObjectId
from datetime import datetime
from .route import Route
from . import Service, id_generator
from asyncio.queues import Queue
from asyncio import StreamWriter, StreamReader
from typing import Optional, Callable, Coroutine, Dict, Generator, Any, List, Union, TypeVar
from dataclasses import dataclass


class RpcEnum(object):
    request = 1  # 请求
    response = 2  # 响应
    notice = 3  # 通知

    @classmethod
    def has(cls, msg_type):
        return msg_type in [
            cls.request,
            cls.response,
            cls.notice,
        ]


Head_Len = struct.calcsize("!ibi")  # 请求头长度  !ibi表示: data_length, msg_type, msg_id


def create_future() -> asyncio.Future:
    return asyncio.get_running_loop().create_future()


def encode(obj):
    if isinstance(obj, ObjectId):
        return {'__ObjectId__': True, '_id': str(obj)}
    elif isinstance(obj, datetime):
        return {'__datetime__': True, 'ttltime': obj.__str__()}
    return obj


def decode(obj):
    if '__ObjectId__' in obj:
        obj = ObjectId(obj['_id'])
    elif '__datetime__' in obj:
        # obj = datetime.strptime(obj['ttltime'], '%Y-%m-%d %H:%M:%S')
        obj = obj['ttltime']

    return obj


class Conn(object):
    """处理写入读取"""

    def __init__(self, reader: StreamReader, writer: StreamWriter, name: str = None, run_send: bool = True):
        self.reader = reader
        self.writer = writer
        self.name = name
        self.queue: Queue = Queue()
        self.send_task: Optional[asyncio.Task] = None

        if run_send:
            self.do_send()

    def do_send(self) -> None:
        """开始数据发送"""
        if self.send_task:
            return

        self.send_task = asyncio.create_task(self._write(), name=self.name)

        def task_done(task):
            self.send_task = None

        self.send_task.add_done_callback(task_done)

    async def _write(self) -> None:
        """队列的形式保证消息发送顺序执行"""
        while True:
            msg = await self.queue.get()
            if msg == "break":
                break

            try:
                self.writer.write(msg)
                await self.writer.drain()
            except ConnectionResetError:
                break
            except Exception as e:
                # print(f"handle._write error=>{e}")
                break

        try:
            self.writer.close()
            await self.writer.wait_closed()
        except ConnectionResetError:
            pass

    def write(self, msg: bytes) -> None:
        self.queue.put_nowait(msg)

    async def read(self) -> bytes:
        return await self.reader.read(1024)

    def close(self) -> None:
        self.queue.put_nowait("break")


MsgId = TypeVar("MsgId", bound="int")


@dataclass
class Request:
    ctime: int  # 创建时间
    result: Union[Callable[..., None], Coroutine, asyncio.Future]  # 回调函数 or Future
    args: Union[list, tuple]  # 回调函数参数
    kwargs: Optional[Dict]  # 回调参数缺省参数


class RpcHandle(object):
    """
    rpc消息统一处理
    """

    def __init__(
            self,
            parent,
            conn: Conn,
            route: Route,
            close_callback: Callable = None,

    ):
        self._data = b""
        self.parent: Service = parent
        self.conn: Conn = conn
        self._id_iter: Generator = id_generator()  # 消息id生成器
        self.route: Route = route  # 服务路由
        self.close_callback: Optional[Callable[[], None]] = close_callback  # 链接关闭回调
        self._request_table: Dict[MsgId, Request] = {}  # 保存所有的rpc请求的回调
        self.timeout: int = 60  # 请求超时时间
        self.log: Log = Log("default")
        self.is_close = False

        # 过期请求消息清理
        asyncio.create_task(self.auto_clear_request())

    def set_log(self, log: Log) -> None:
        """设置日志容器"""
        self.log = log

    async def auto_clear_request(self):
        """清除超时的请求"""
        while not self.is_close:
            await asyncio.sleep(self.timeout + 5)
            del_list = []
            nt = int(time.time())
            for msg_id, request in self._request_table.items():
                if nt - request.ctime > self.timeout:
                    del_list.append(msg_id)

            if del_list:
                for msg_id in del_list:
                    self._request_table.pop(msg_id)

    async def run(self) -> None:
        while True:
            try:
                msg = await self.conn.read()
            except ConnectionResetError:
                break
            except Exception as e:
                print(f"RpcHandle.run error=>{e}")
                break

            if msg == b"":
                # print(f"对端断开")
                break
            elif msg == b"stop_server":
                await self.parent.close()
                continue

            await self.handle_msg(msg)

        self.is_close = True
        self.conn.close()
        if callable(self.close_callback):
            await self.close_callback(self)

    async def handle_msg(self, msg: bytes) -> None:
        """粘包拆包 消息处理"""
        self._data += msg

        cur_data_len = len(self._data)
        while cur_data_len > Head_Len:
            body_length, msg_type, msg_id = struct.unpack('!ibi', self._data[:Head_Len])
            if not RpcEnum.has(msg_type):
                self._data = b""
                self.log.warning(f"非法请求 msg_id:{msg_id}")
                break

            # 数据不完整，继续接收
            if cur_data_len - Head_Len < body_length:
                break

            body = self._data[Head_Len:Head_Len + body_length]
            self._data = self._data[Head_Len + body_length:]
            cur_data_len = len(self._data)
            try:
                body = msgpack.loads(body, object_hook=decode, strict_map_key=False)
            except Exception as e:
                print(f"msgpack.loads失败[{self.parent.name}] err: {e}")
                break

            if msg_type == RpcEnum.request:
                asyncio.create_task(self.handle_request(msg_id, body))
            elif msg_type == RpcEnum.response:
                asyncio.create_task(self.handle_response(msg_id, body))
            elif msg_type == RpcEnum.notice:
                asyncio.create_task(self.handle_notice(body))

    @staticmethod
    @capture_error()
    def _pack_request(msg_id: int, msg_type: int, method_name: str, *args, **kwargs) -> bytes:
        """打包请求"""
        # 消息体
        body = msgpack.dumps((method_name, args, kwargs), default=encode)
        # 请求头
        head = struct.pack('!ibi', len(body), msg_type, msg_id)

        return head + body

    @staticmethod
    @capture_error()
    def _pack_response(msg_id: int, msg_type: int, response: dict, err: str) -> bytes:
        """打包响应"""
        # 消息体
        body = msgpack.dumps((response, err), default=encode)
        # 请求头
        head = struct.pack('!ibi', len(body), msg_type, msg_id)

        return head + body

    async def handle_request(self, msg_id: int, body: tuple) -> None:
        """处理请求"""

        method_name, args, kwargs = body

        if method_name == "register":
            args.insert(0, self)

        response, err = None, None
        try:
            response = await self.route.call(method_name, *args, **kwargs)
        except Exception:
            err = str(traceback.format_exc())
            self.log.error(f"rpc服务器有报错 handle_request[{method_name}] args=>{args} kwargs=>{kwargs} {err}")

        _data = self._pack_response(msg_id, RpcEnum.response, response, err)
        self.conn.write(_data)

    async def handle_response(self, msg_id: int, body: tuple) -> None:
        """处理响应"""
        request = self.get_request(msg_id)
        # time out?
        if not request:
            return

        response, err = body

        result = request.result
        # 处理call方法的future,设置结果
        if isinstance(result, asyncio.Future):
            result.set_result(response)
        # 处理callback方法的回调
        elif callable(result):
            result(response, err, *request.args, **request.kwargs)

    async def handle_notice(self, body) -> None:
        """处理通知"""
        method_name, args, kwargs = body
        try:
            await self.route.call(method_name, *args, **kwargs)
        except:
            err = str(traceback.format_exc())
            self.log.error(f"rpc服务器有报错 handle_notice[{method_name}] args=>{args} kwargs=>{kwargs} {err}")

    def add_request(self, msg_id: int, result: Union[asyncio.Future, Callable], *args, **kwargs) -> None:
        """记录请求的消息"""
        self._request_table[msg_id] = Request(
            ctime=int(time.time()),
            result=result,
            args=args,
            kwargs=kwargs,
        )

    def get_request(self, msg_id: int, pop: bool = True) -> Optional[Request]:
        """获取请求消息"""
        if msg_id not in self._request_table:
            return

        if pop:
            return self._request_table.pop(msg_id)
        else:
            return self._request_table[msg_id]

    def notice(self, func_name: str, *args, **kwargs) -> None:
        """非阻塞 通知"""
        _data = self._pack_request(next(self._id_iter), RpcEnum.notice, func_name, *args, **kwargs)
        if not _data:
            return

        self.conn.write(_data)

    def callback(self, func_name: str, func_args: List[Any], callback_func: Callable, *args, **kwargs) -> None:
        """非阻塞 回调 args的最后一个参数为回调函数"""
        msg_id = next(self._id_iter)
        data = self._pack_request(msg_id, RpcEnum.request, func_name, *func_args)
        if not data:
            return

        self.add_request(msg_id, callback_func, *args, **kwargs)
        self.conn.write(data)

    async def call(self, func_name: str, *args, **kwargs) -> Any:
        """阻塞 等待结果返回"""
        msg_id = next(self._id_iter)

        data = self._pack_request(msg_id, RpcEnum.request, func_name, *args, **kwargs)
        if not data:
            return

        call_future = create_future()
        self.add_request(msg_id, call_future)
        self.conn.write(data)

        try:
            return await asyncio.wait_for(call_future, self.timeout)
        except TimeoutError:
            err = traceback.format_exc()
            self.log.warning(f"handle.call {func_name} timeout args:{args}, kwargs:{kwargs} err=>{err}")
        except Exception as e:
            print(f"handle.call error=>{e}")


if __name__ == '__main__':
    pass
